import uuid
from flask import Blueprint, request, jsonify
from services.chroma_service import ChromaService
from models.document import Knowledge
from utils.tools import deterministic_uuid
import json

vector_bp = Blueprint('vector', __name__)
chroma_service = ChromaService()

#######################################
# API：Collection 的增删改查
#######################################

# 查询所有 collections
@vector_bp.route('/', methods=['GET'])
def list_collections():
    # 获取分页和搜索参数
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    search = request.args.get("search", "").lower()
    
    try:
        result = chroma_service.list_collections(page, page_size, search)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 新建 collection
@vector_bp.route('/', methods=['POST'])
def create_collection():
    data = request.get_json()
    name = data.get('name')
    if not name:
        return jsonify({"error": "Collection name is required"}), 400
    try:
        result = chroma_service.create_collection(name)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    
# 修改 collection（例如重命名操作：复制数据到新集合然后删除旧集合）
@vector_bp.route('/<collection_name>', methods=['PUT'])
def update_collection(collection_name):
    data = request.get_json()
    new_name = data.get('new_name')
    if not new_name:
        return jsonify({"error": "New collection name is required"}), 400
    try:
        result = chroma_service.update_collection(collection_name, new_name)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 删除 collection
@vector_bp.route('/<collection_name>', methods=['DELETE'])
def delete_collection(collection_name):
    try:
        result = chroma_service.delete_collection(collection_name)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    
#######################################
# API：记录的增删改查与检索（在具体 collection 下）
#######################################

# 分页查询记录，支持简单的“搜索”（在 document 字段中模糊匹配）
@vector_bp.route('/<collection_name>/records', methods=['GET'])
def list_records(collection_name):
    # 获取分页参数
    page = int(request.args.get("page", 1))
    page_size = int(request.args.get("page_size", 10))
    
    # 获取查询参数
    search_type = request.args.get("search_type")
    record_id = request.args.get("id")
    query_text = request.args.get("query_text")
    metadata_filter = request.args.get("metadata_filter")
    document_filter = request.args.get("document_filter")
    n_results_requested = int(request.args.get("n_results", 1000))
    
    try:
        # 调用 ChromaService 的 list_records 方法
        result = chroma_service.list_records(
            collection_name=collection_name,
            embedding_model=None,
            page=page,
            page_size=page_size,
            search_type=search_type,
            record_id=record_id,
            query_text=query_text,
            metadata_filter=metadata_filter,
            document_filter=document_filter,
            n_results_requested=n_results_requested
        )
        return jsonify(result)
    except Exception as e:
        print(f"Error in list_records: {str(e)}")  # 添加错误日志
        return jsonify({"error": str(e)}), 500

# 新建记录
@vector_bp.route('/<collection_name>/records', methods=['POST'])
def create_record(collection_name):
    data = request.get_json()
    record_id = data.get("id", str(uuid.uuid4()))
    document = data.get("document", "")
    # 如果 metadata 不存在或为空，则不包含在参数中
    metadata = data.get("metadata")
    
    if not document:
        return jsonify({"error": "Document content is required"}), 400
        
    try:
        result = chroma_service.create_record(collection_name, record_id, document, metadata)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500


# 更新记录（模拟：先删除再添加）
@vector_bp.route('/<collection_name>/records/<record_id>', methods=['PUT'])
def update_record(collection_name, record_id):
    data = request.get_json()
    document = data.get("document", "")
    # 如果 metadata 不存在或为空，则不包含在参数中
    metadata = data.get("metadata")
    
    if not document:
        return jsonify({"error": "Document content is required"}), 400
        
    try:
        result = chroma_service.update_record(collection_name, record_id, document, metadata)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 删除单个记录
@vector_bp.route('/<collection_name>/records/<record_id>', methods=['DELETE'])
def delete_record(collection_name, record_id):
    try:
        result = chroma_service.delete_record(collection_name, record_id)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500
    
# 批量删除记录，传入 JSON 格式的 { "ids": [id1, id2, ...] }
@vector_bp.route('/<collection_name>/records', methods=['DELETE'])
def batch_delete_records(collection_name):
    data = request.get_json()
    ids = data.get("ids")
    if not ids or not isinstance(ids, list):
        return jsonify({"error": "List of record ids is required"}), 400
    try:
        result = chroma_service.batch_delete_records(collection_name, ids)
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

#######################################
# API：批量保存向量数据
#######################################

@vector_bp.route('/batch-save-structured', methods=['POST'])
def batch_save_structured():
    data = request.get_json()
    
    # 参数验证
    if not data.get('knowledge_id'):
        return jsonify({"error": "knowledge_id is required"}), 400
    
    if not data.get('records') or not isinstance(data.get('records'), list):
        return jsonify({"error": "records array is required"}), 400
        
    if not data.get('document_fields'):
        return jsonify({"error": "document_fields is required"}), 400
        
    # 获取knowledge信息
    knowledge = Knowledge.query.get(data['knowledge_id'])
    if not knowledge:
        return jsonify({"error": "Knowledge not found"}), 404
    collection_name = knowledge.collection_name
    embedding_model = knowledge.embedding_model
        
    use_content_embedding = data.get('use_content_embedding', False)  # 重命名字段
    
    if not use_content_embedding and not data.get('embedding_fields'):
        return jsonify({"error": "embedding_fields is required when use_content_embedding is false"}), 400
    
    try:
        # 处理数据
        records = data['records']
        
        # 构建documents
        documents = []
        for record in records:
            doc_content = {field: str(record.get(field, '')) for field in data['document_fields']}
            documents.append(json.dumps(doc_content, ensure_ascii=False))

        # 使用传入的 ids 或生成新的
        provided_ids = data.get('ids')
        ids = provided_ids or [deterministic_uuid(doc) for doc in documents]
        if len(ids) != len(documents):
            return jsonify({"error": "ids length must match documents length"}), 400
            
        # 构建metadata
        metadatas = []
        for record in records:
            metadata = {
                field: record.get(field) 
                for field in data.get('metadata_fields', []) 
                if field in record
            }
            metadata['knowledge_id'] = data['knowledge_id']
            metadatas.append(metadata)
            
        # 获取embeddings
        chroma_service = ChromaService()
        if use_content_embedding:
            embeddings = [chroma_service.get_embedding(doc, embedding_model) for doc in documents]
        else:
            embed_contents = []
            for record in records:
                embed_content = {field: str(record.get(field, '')) for field in data['embedding_fields']}
                embed_contents.append(json.dumps(embed_content, ensure_ascii=False))
            embeddings = [chroma_service.get_embedding(content, embedding_model) for content in embed_contents]
            
        # 保存数据
        result = chroma_service.add_records(
            collection_name=collection_name,
            embeddings=embeddings,
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )
        
        return jsonify({
            "success": True,
            "message": "Data saved successfully",
            "result": result,
            "ids": ids if not provided_ids else None  # 只在未提供 ids 时返回生成的 ids
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@vector_bp.route('/batch-save-raw', methods=['POST'])
def batch_save_raw():
    data = request.get_json()
    
    # 参数验证
    if not data.get('knowledge_id'):
        return jsonify({"error": "knowledge_id is required"}), 400
        
    documents = data.get('document', [])
    metadatas = data.get('metadata', [])
    embedding_contents = data.get('embedding', [])
    use_content_embedding = data.get('use_content_embedding', False)
    
    # 验证数据长度一致性
    if not documents:
        return jsonify({"error": "documents array is required"}), 400
        
    if metadatas and len(metadatas) != len(documents):
        return jsonify({"error": "metadata length must match documents length"}), 400
        
    if not use_content_embedding and (not embedding_contents or len(embedding_contents) != len(documents)):
        return jsonify({"error": "embedding contents length must match documents length"}), 400
    
    # 获取knowledge信息
    knowledge = Knowledge.query.get(data['knowledge_id'])
    if not knowledge:
        return jsonify({"error": "Knowledge not found"}), 404
    collection_name = knowledge.collection_name
    embedding_model = knowledge.embedding_model
    
    try:
        # 使用传入的 ids 或生成新的
        provided_ids = data.get('ids')
        ids = provided_ids or [deterministic_uuid(doc) for doc in documents]
        if len(ids) != len(documents):
            return jsonify({"error": "ids length must match documents length"}), 400
            
        # 处理metadata
        processed_metadatas = []
        for meta in (metadatas or [{}] * len(documents)):
            metadata = meta.copy() if meta else {}
            metadata['knowledge_id'] = data['knowledge_id']
            processed_metadatas.append(metadata)
            
        # 获取embeddings
        chroma_service = ChromaService()
        if use_content_embedding:
            embeddings = [chroma_service.get_embedding(doc, embedding_model) for doc in documents]
        else:
            embeddings = [chroma_service.get_embedding(content, embedding_model) for content in embedding_contents]
            
        # 保存数据
        result = chroma_service.add_records(
            collection_name=collection_name,
            embeddings=embeddings,
            documents=documents,
            metadatas=processed_metadatas,
            ids=ids
        )
        
        return jsonify({
            "success": True,
            "message": "Data saved successfully",
            "result": result,
            "ids": ids if not provided_ids else None  # 只在未提供 ids 时返回生成的 ids
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@vector_bp.route('/batch-delete', methods=['POST'])
def batch_delete():
    data = request.get_json()
    
    # 参数验证
    if not data.get('knowledge_id'):
        return jsonify({"error": "knowledge_id is required"}), 400
        
    if not data.get('ids') or not isinstance(data.get('ids'), list):
        return jsonify({"error": "ids array is required"}), 400
    
    # 获取knowledge信息
    knowledge = Knowledge.query.get(data['knowledge_id'])
    if not knowledge:
        return jsonify({"error": "Knowledge not found"}), 404
    collection_name = knowledge.collection_name
    
    try:
        # 删除数据
        result = chroma_service.batch_delete_records(
            collection_name=collection_name,
            ids=data['ids']
        )
        
        return jsonify({
            "success": True,
            "message": "Data deleted successfully",
            "result": result
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500